#!/usr/bin/env python3
"""
从标注点生成分割掩码
支持心内膜、心外膜和心肌掩码生成
"""

import numpy as np
import cv2
from typing import Dict, List, Tuple, Optional
from pathlib import Path
from src.utils.logger import get_logger

logger = get_logger("MaskGenerator")


class MaskGenerator:
    """从轮廓点生成分割掩码"""
    
    def __init__(self, image_shape: Tuple[int, int] = (112, 112)):
        """
        初始化掩码生成器
        
        Args:
            image_shape: 图像形状 (height, width)
        """
        self.image_shape = image_shape
        self.height, self.width = image_shape
    
    def _normalize_points(self, points: np.ndarray, 
                          source_shape: Optional[Tuple[int, int]] = None) -> np.ndarray:
        """
        归一化坐标点到目标图像尺寸
        
        Args:
            points: 轮廓点 (N, 2)
            source_shape: 源图像尺寸 (height, width)，如果为None则假设坐标已经是归一化的
            
        Returns:
            归一化后的点 (N, 2)
        """
        if points is None or len(points) == 0:
            return points
        
        points = points.copy()
        
        if source_shape is not None:
            # 从源尺寸归一化
            src_h, src_w = source_shape
            points[:, 0] = points[:, 0] * self.width / src_w
            points[:, 1] = points[:, 1] * self.height / src_h
        else:
            # 假设坐标已经是像素坐标，但可能需要缩放
            # 检查坐标范围，如果在0-120左右，可能需要缩放到图像尺寸
            if points[:, 0].max() > self.width or points[:, 1].max() > self.height:
                # 需要缩放
                scale_x = self.width / (points[:, 0].max() - points[:, 0].min() + 1)
                scale_y = self.height / (points[:, 1].max() - points[:, 1].min() + 1)
                scale = min(scale_x, scale_y) * 0.9  # 留一些边距
                
                # 先平移到原点
                points[:, 0] -= points[:, 0].min()
                points[:, 1] -= points[:, 1].min()
                
                # 缩放
                points[:, 0] *= scale
                points[:, 1] *= scale
                
                # 平移到中心
                points[:, 0] += (self.width - points[:, 0].max()) / 2
                points[:, 1] += (self.height - points[:, 1].max()) / 2
        
        # 裁剪到图像范围内
        points[:, 0] = np.clip(points[:, 0], 0, self.width - 1)
        points[:, 1] = np.clip(points[:, 1], 0, self.height - 1)
        
        return points
    
    def generate_lv_mask(self, 
                        endo_points: Optional[np.ndarray] = None,
                        epi_points: Optional[np.ndarray] = None,
                        mask_type: str = 'binary') -> np.ndarray:
        """
        生成左心室掩码
        
        Args:
            endo_points: 心内膜轮廓点 (N, 2)
            epi_points: 心外膜轮廓点 (M, 2)
            mask_type: 掩码类型
                - 'binary': 二值掩码（左心室区域=255）
                - 'multi_class': 多类别掩码（心内膜=1，心肌=2，心外膜=3）
        
        Returns:
            掩码数组 (height, width)
        """
        if mask_type == 'binary':
            return self.generate_binary_mask(endo_points)
        elif mask_type == 'multi_class':
            return self.generate_multi_class_mask(endo_points, epi_points)
        else:
            raise ValueError(f"不支持的掩码类型: {mask_type}")
    
    def generate_binary_mask(self, contour_points: Optional[np.ndarray]) -> np.ndarray:
        """
        生成二值掩码（左心室区域）
        
        Args:
            contour_points: 轮廓点 (N, 2)
            
        Returns:
            二值掩码 (height, width)，左心室区域=255，背景=0
        """
        mask = np.zeros(self.image_shape, dtype=np.uint8)
        
        if contour_points is None or len(contour_points) < 3:
            return mask
        
        # 归一化点
        points = self._normalize_points(contour_points)
        
        # 转换为整数坐标
        points_int = points.astype(np.int32)
        
        # 填充轮廓
        if len(points_int) >= 3:
            cv2.fillPoly(mask, [points_int], 255)
        
        return mask
    
    def generate_endo_mask(self, endo_points: Optional[np.ndarray]) -> np.ndarray:
        """
        生成心内膜掩码
        
        Args:
            endo_points: 心内膜轮廓点 (N, 2)
            
        Returns:
            心内膜掩码 (height, width)
        """
        return self.generate_binary_mask(endo_points)
    
    def generate_epi_mask(self, epi_points: Optional[np.ndarray]) -> np.ndarray:
        """
        生成心外膜掩码
        
        Args:
            epi_points: 心外膜轮廓点 (M, 2)
            
        Returns:
            心外膜掩码 (height, width)
        """
        return self.generate_binary_mask(epi_points)
    
    def generate_myocardium_mask(self, 
                                 endo_points: Optional[np.ndarray],
                                 epi_points: Optional[np.ndarray]) -> np.ndarray:
        """
        生成心肌掩码（心外膜 - 心内膜）
        
        Args:
            endo_points: 心内膜轮廓点 (N, 2)
            epi_points: 心外膜轮廓点 (M, 2)
            
        Returns:
            心肌掩码 (height, width)，心肌区域=255，其他=0
        """
        if epi_points is None or endo_points is None:
            return np.zeros(self.image_shape, dtype=np.uint8)
        
        # 生成心外膜和心内膜掩码
        epi_mask = self.generate_epi_mask(epi_points)
        endo_mask = self.generate_endo_mask(endo_points)
        
        # 心肌 = 心外膜 - 心内膜
        myocardium_mask = cv2.subtract(epi_mask, endo_mask)
        
        return myocardium_mask
    
    def generate_multi_class_mask(self,
                                  endo_points: Optional[np.ndarray],
                                  epi_points: Optional[np.ndarray]) -> np.ndarray:
        """
        生成多类别掩码
        
        Args:
            endo_points: 心内膜轮廓点 (N, 2)
            epi_points: 心外膜轮廓点 (M, 2)
            
        Returns:
            多类别掩码 (height, width)
                - 0: 背景
                - 1: 心内膜区域
                - 2: 心肌区域
                - 3: 心外膜区域（如果只有心内膜，则为心内膜区域）
        """
        mask = np.zeros(self.image_shape, dtype=np.uint8)
        
        if endo_points is None:
            return mask
        
        # 生成心内膜掩码（类别1）
        endo_mask = self.generate_endo_mask(endo_points)
        mask[endo_mask > 0] = 1
        
        if epi_points is not None:
            # 生成心外膜掩码（类别3）
            epi_mask = self.generate_epi_mask(epi_points)
            mask[epi_mask > 0] = 3
            
            # 生成心肌掩码（类别2）
            myocardium_mask = self.generate_myocardium_mask(endo_points, epi_points)
            mask[myocardium_mask > 0] = 2
        
        return mask
    
    def generate_mask_from_contours(self,
                                    contours: Dict[str, Optional[np.ndarray]],
                                    mask_type: str = 'binary') -> np.ndarray:
        """
        从轮廓字典生成掩码
        
        Args:
            contours: 轮廓字典，包含 'endo' 和/或 'epi' 键
            mask_type: 掩码类型
            
        Returns:
            掩码数组
        """
        endo_points = contours.get('endo')
        epi_points = contours.get('epi')
        
        return self.generate_lv_mask(endo_points, epi_points, mask_type)


def create_mask_from_annotation(video_id: str,
                               frame_num: int,
                               annotation_data: Dict,
                               image_shape: Tuple[int, int] = (112, 112),
                               mask_type: str = 'binary') -> Optional[np.ndarray]:
    """
    从标注数据创建掩码（便捷函数）
    
    Args:
        video_id: 视频ID
        frame_num: 帧号
        annotation_data: 标注数据字典，包含 'endo' 和/或 'epi' 轮廓点
        image_shape: 图像形状
        mask_type: 掩码类型
        
    Returns:
        掩码数组或None
    """
    generator = MaskGenerator(image_shape)
    
    endo_points = annotation_data.get('endo')
    epi_points = annotation_data.get('epi')
    
    if endo_points is None and epi_points is None:
        return None
    
    return generator.generate_lv_mask(endo_points, epi_points, mask_type)

